Supercharge xpx.testing.lazy_xp_function.

This is a fairly high level change; @rgommers @pearu @jakevdp I could use your feedback.

Matching SciPy PR: scipy/scipy#22909

Automatic static arguments

scipy has seen a proliferation of lazy_xp_function(func, static_argnames=(...)) in the initial section of many test modules. This information is only useful for JAX and is quite verbose.
With scipy/scipy#22686, this informational quirk that is both specific to JAX and to unit tests moves to the implementation modules.

With this PR, lazy_xp_function no longer accepts parameters static_argnames and static_argnums. Instead, all arguments that are not JAX arrays are automatically treated as static. Note that this behaviour is generally desirable in unit testing but a bad idea in production. Consider:

def f(x: Array, y: float, plus: bool) -> Array:
    return x + y if plus else x - y

j1 = jax.jit(f, static_argnames="plus")
j2 = jax_autojit(f)

jax_autojit is a new internal function of array-api-extra.
In the above example, j2 requires a lot less setup to be tested effectively, but on the flip side it means that it will be re-traced for every different value of y, which likely makes it not fit for purpose in production.

To clarify: jax_autojit is applied and removed on the fly within tests with the xp fixture and it is never used outside of unit testing.

Wrapped inputs and output

There are a few cases of scipy functions returning bespoke containers with arrays inside instead of simple tuples, namedtuples, lists, or dicts of arrays.

Two such examples are

• scipy.stats.ttest_ind
• scipy.spatial.Rotation and all functions and methods that return it (discussion: ENH: spatial.transform: add array API standard support scipy/scipy#22777 (comment))

In main, these can't be tested with lazy_xp_function, and as of scipy/scipy#22686 the issue also impacts documentation.

This PR lifts this restriction and allows completely arbitrary objects as parameters and as return values of the functions. If these objects internally contain JAX arrays, lazy_xp_function will now automatically extract them, pass them through the JIT, and reassemble everything for the test function to observe the result.

The rationale is that, in real life, users are unlikely to wrap the scipy functions with jax.jit directly; instead they are more likely to consume their outputs in their own functions and then wrap those with jax.jit:

@jax.jit
def my_user_function(x):
    y = scipy.stats.ttest_ind(x)
    return my_user_consume(y)  # returns array or tuple of arrays

Static non-hashable arguments

Non-hashable objects can now be static.
Consider:

def f(x: Array, verbs: Sequence[str]) -> Array:
    ...
    
lazy_xp_function(f, static_argnames="verbs")

def test_f(xp):
    x = ...
    y = f(x, ("foo", "bar", "baz"))  # OK
    y = f(x, ["foo", "bar", "baz"])  # Fails

This PR fixes it; all objects in the input and output of the function need only be hashable or pickleable.

Dask materialization raises inside a container

This also fixes a Dask-specific bug where the graph materialization raises:

@pytest.mark.skip_xp_backends("jax.numpy", reason="raise inside jax.pure_callback")
def test1(xp):
    with pytest.raises(Exception):
        f(x)

The above works when f returns plain Dask array objects or tuples or lists thereof, even if the graph would otherwise be discarded and never computed, thanks to the lazy_xp_function machinery, but used to fail when the return value is an opaque container with Dask arrays inside. This PR fixes it.